package shop.milvus.ServiceImpl;

import cn.hutool.core.lang.Assert;
import io.milvus.client.MilvusClient;
import io.milvus.grpc.*;
import io.milvus.param.R;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;
import lombok.AllArgsConstructor;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import shop.milvus.MilvusService;
import shop.milvus.config.MilvusProperties;
import shop.milvus.dto.MilvusSearchDto;

import java.util.*;
import java.util.stream.Collectors;

@Service
@Slf4j
@Deprecated
public class MilvusServiceImpl implements MilvusService {
    @Autowired
    private  MilvusProperties milvusProperties;

    public  MilvusClient client;

    public MilvusClient getClient(){
        return client;
    }
    public MilvusServiceImpl(MilvusClient client) {
        this.client = client;

    }

    @Override
    @Deprecated
    public void createCollection() {
        String collectionName = milvusProperties.getMilvusCollectionName();

        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .build();
        FieldType goodsField = FieldType.newBuilder()
                .withName("goods_id")
                .withDataType(DataType.Int64)
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(128) // 向量维度
                .withAutoID(true)
                .build();

        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("Example collection for Java")
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(vectorField)
                .build();

        client.createCollection(createCollectionReq);
        System.out.println("Collection created!");
    }

    @Override
    public Long insertData(String emdeddingName, List<Float> vector, HashMap<String, Object> fields) {

        ArrayList<InsertParam.Field> insertFields = new ArrayList<>();

        insertFields.add(new InsertParam.Field(emdeddingName,Collections.singletonList(vector)));
        //遍历插入可能的字段键和值
        for (Map.Entry<String, Object> stringObjectEntry : fields.entrySet()) {
            insertFields.add(new InsertParam.Field(stringObjectEntry.getKey(),Collections.singletonList(stringObjectEntry.getValue())));
        }
        //执行插入
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(milvusProperties.getMilvusCollectionName())
                .withFields(insertFields)
                .build();

        R<MutationResult> insertResult = client.insert(insertParam);

        if (insertResult.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("Insert failed: " + insertResult.getMessage());
        }

        Long inserted  = insertResult.getData().getInsertCnt();


        Assert.isTrue(inserted==1,()->new RuntimeException("插入失败"));
        return inserted;
    }

    @Override
    public List<MilvusSearchDto> searchVectors(List<Float> vector, String fieldName, int topK, List<String> outFieldNames, String toSoleFieldName, List<Long> skipIds, String filterExpr) {

        SearchParam.Builder builder = SearchParam.newBuilder()
                .withCollectionName(milvusProperties.getMilvusCollectionName())
                .withVectorFieldName(fieldName)
                .withTopK(topK*milvusProperties.getRedundancy())
                .withFloatVectors(Collections.singletonList(vector))
                .withOutFields(outFieldNames)
                .withMetricType(milvusProperties.getMetricType())
                .withParams("{\"nprobe\":"+milvusProperties.getTemperature()+"}");
        //超级拼装
        List<String> exprParts = new ArrayList<>();
        //跳过的id集合
        if (skipIds != null && !skipIds.isEmpty()) {
            String skipExpr = "NOT id in [" + skipIds.stream()
                    .map(String::valueOf)
                    .collect(Collectors.joining(",")) + "]";
            exprParts.add(skipExpr);
        }
        //过滤条件
        if (filterExpr != null && !filterExpr.isEmpty()) {
            exprParts.add("(" + filterExpr + ")");
        }

        String expr = String.join(" and ", exprParts);

        if (!expr.isEmpty()) {
            builder.withExpr(expr);
        }
           SearchParam searchParam= builder.build();

        R<SearchResults> results = client.search(searchParam);
        SearchResults data = results.getData();
        SearchResultData results1 = data.getResults();
        SearchResultsWrapper wrapper = new SearchResultsWrapper(results1);

        List<SearchResultsWrapper.IDScore> idScores = wrapper.getIDScore(0);

        List<MilvusSearchDto> result = new ArrayList<>();
        //去重集合
        HashSet<Object> sole = new HashSet<>((int)(topK / 0.75f) + 1);

        // 遍历 topK 结果（RowRecord 和 IDScore 一一对应）
        for (int i = 0; i < wrapper.getRowRecords().size(); i++) {


            QueryResultsWrapper.RowRecord rowRecord = wrapper.getRowRecords().get(i);
            SearchResultsWrapper.IDScore idScore = idScores.get(i);
            if (!sole.contains(rowRecord.get(toSoleFieldName))&&sole.size()<topK) {

                //加入去重集合
                sole.add(rowRecord.get(toSoleFieldName));

                MilvusSearchDto dto = new MilvusSearchDto();
                dto.setId(idScore.getLongID());
                dto.setScore(idScore.getScore());


                // 额外字段
                Map<String, Object> extraFields = new HashMap<>();
                for (String outFieldName : outFieldNames) {
                    Object fieldValue = rowRecord.get(outFieldName);
                    extraFields.put(outFieldName, fieldValue);
                    log.info("{}={}", outFieldName, fieldValue);
                }
                dto.setExtraFields(extraFields);

                log.info("id={}, score={}", idScore.getLongID(), idScore.getScore());

                result.add(dto);
            }
        }
        return result;
    }


    @Override
    public long deleteDataById( List<Long> ids) {
        if (ids == null || ids.isEmpty()) return 0;

        String expr = "id in " + ids.toString();

        DeleteParam deleteParam = DeleteParam.newBuilder()
                .withCollectionName(milvusProperties.getMilvusCollectionName())
                .withExpr(expr)
                .build();

        R<MutationResult> response = client.delete(deleteParam);

        if (response.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("删除失败: " + response.getMessage());
        }

        // 返回删除的行数
        return response.getData().getDeleteCnt();
    }
    @Override
    public long deleteDataByGoodsId(  List<Long> ids) {
        if (ids == null || ids.isEmpty()) return 0;

        String expr = "id in " + ids.toString();

        DeleteParam deleteParam = DeleteParam.newBuilder()
                .withCollectionName(milvusProperties.getMilvusCollectionName())
                .withExpr(expr)
                .build();

        R<MutationResult> response = client.delete(deleteParam);

        if (response.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("删除失败: " + response.getMessage());
        }

        // 返回删除的行数
        return response.getData().getDeleteCnt();
    }

}
